from advhat import Advhat
import skimage.io as io
from utils import *
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import cv2
import numpy as np
device = "cpu"
if __name__ == "__main__":
    host_dir = "./before_aligned_600" #需要修改的图像
    target_dir = "./target_aligned_600/Camilla_Parker_Bowles_0002_aligned.png"#攻击的目标
    target_img = io.imread(target_dir)/255.0

    # 额头贴纸的样式图
    logo_dir = "./logo/example.png"
    logo_img = io.imread(logo_dir)/255.0
    # logo_img = np.ones((400, 900, 3)) * 127. / 255.
    # numpy转tensor
    transform = transforms.Compose([transforms.ToTensor(),transforms.Resize([600,600])])
    transform1 = transforms.Compose([transforms.ToTensor()])

    file_list = os.listdir(host_dir)
    model = Advhat(1).to(dtype=torch.float32,device=device) # 获取模型
    target_img = torch.from_numpy(target_img).unsqueeze(0).to(device=device, dtype=torch.float32).permute(0, 3, 1, 2)
    for file in file_list:
        path = os.path.join(host_dir, file)
        host_img = io.imread(path)
        host_img = torch.from_numpy(host_img).to(device=device,dtype=torch.float32)
        logo_tensor = torch.from_numpy(logo_img).to(device=device,dtype=torch.float32).unsqueeze(0).permute(0,3,1,2)

        # 进入模型的训练阶段
        """
        state1需要iter==100，state2需要iter2==200
        """
        logo,moments,host = model.FGSM_with_moments(host_img,target_img,logo_tensor,200)
        moments = moments.to(device)
        logo,final_result = model.FGSM_with_moments(host_img,target_img,logo,400,moments)
        name = file.split(".")

        logo_numpy = logo.cpu().permute(1,2,0).detach().numpy().astype(np.float32)*255
        final_numpy = final_result.permute(1,2,0).cpu().detach().numpy().astype(np.float32)*255

        # 最终可以保存下来的对抗结果

        logo_numpy = cv2.cvtColor(logo_numpy,cv2.COLOR_RGB2BGR).astype(np.uint8)# 这是贴图
        final_numpy = cv2.cvtColor(final_numpy, cv2.COLOR_RGB2BGR).astype(np.uint8)# 这是贴图与人物

        cv2.imwrite("./face_with_logo/" + name[0] + ".png",  final_numpy)
        cv2.imwrite("./face_with_logo/" + 'logo' + ".png", logo_numpy)
